import pickle
import faiss
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import torch
import torchvision.transforms as T
from PIL import Image

patch_h = 28
patch_w = 28
feat_dim = 384

transform = T.Compose([
    T.GaussianBlur(9, sigma=(0.1, 2.0)),
    T.Resize((patch_h * 14, patch_w * 14)),
    T.CenterCrop((patch_h * 14, patch_w * 14)),
    T.ToTensor(),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])

image_path = './images/dandelion.jpg'

dinov2 = torch.hub.load('./model', 'dinov2_vits14', source='local').cuda()

image = Image.open(image_path)
t_image = transform(image).unsqueeze(dim=0)
querry_feature = dinov2(t_image.cuda())
normalized_feature = querry_feature / querry_feature.norm(dim=-1, keepdim=True)
normalized_feature = normalized_feature.detach().cpu().numpy()

# 加载保存的特征矩阵和URL映射表
with open('./faiss/feature_matrix.pkl', 'rb') as f:
    feature_matrix = pickle.load(f)

with open('./faiss/url_mapping.pkl', 'rb') as f:
    url_mapping = pickle.load(f)

# 创建Faiss索引
index = faiss.IndexFlatIP(feature_matrix.shape[1])
index.add(feature_matrix)

# 执行查询
k = 5  # 要检索的最相似图片数量
D, I = index.search(normalized_feature, k)

# 获取最相似特征向量对应的图片URL
similar_images = []
print("最相似图片的URL：")
for i in range(k):
    similar_image_index = I[0][i]
    similar_image_url = url_mapping[similar_image_index]
    print(similar_image_url)
    similar_images.append(similar_image_url)